# gomoku_ai/ai/classic_ai.py
import numpy as np
import random


class ClassicAI:
    def __init__(self, board_size, difficulty=3):
        self.board_size = board_size
        self.difficulty = difficulty
        self.directions = [(1, 0), (0, 1), (1, 1), (1, -1)]

        # 根据难度调整搜索深度
        if difficulty == 1:
            self.search_depth = 1
        elif difficulty == 2:
            self.search_depth = 2
        elif difficulty == 3:
            self.search_depth = 3
        elif difficulty == 4:
            self.search_depth = 4
        elif difficulty == 5:
            self.search_depth = 5
        else:
            self.search_depth = 3  # 默认

    def get_action(self, board):
        """获取AI的动作"""
        valid_moves = self.get_valid_moves(board)

        if not valid_moves:
            return None

        if self.difficulty == 1:
            return random.choice(valid_moves)
        elif self.difficulty == 2:
            return self.get_medium_action(board, valid_moves)
        else:
            return self.get_best_action(board, valid_moves)

    def get_valid_moves(self, board):
        """获取有效移动"""
        valid_moves = []
        for i in range(self.board_size):
            for j in range(self.board_size):
                if board[i][j] == 0:
                    valid_moves.append(i * self.board_size + j)
        return valid_moves

    def get_medium_action(self, board, valid_moves):
        """中等难度AI"""
        # 首先检查是否能赢
        for action in valid_moves:
            row, col = divmod(action, self.board_size)
            board[row][col] = -1
            if self.check_win(board, -1):
                board[row][col] = 0
                return action
            board[row][col] = 0

        # 然后检查是否需要防守
        for action in valid_moves:
            row, col = divmod(action, self.board_size)
            board[row][col] = 1
            if self.check_win(board, 1):
                board[row][col] = 0
                return action
            board[row][col] = 0

        # 否则随机选择
        return random.choice(valid_moves)

    def get_best_action(self, board, valid_moves):
        """高难度AI"""
        best_score = -float('inf')
        best_action = None

        for action in valid_moves:
            row, col = divmod(action, self.board_size)
            score = self.evaluate_position(board, row, col)

            if score > best_score:
                best_score = score
                best_action = action

        return best_action

    def evaluate_position(self, board, row, col):
        """评估位置分数"""
        score = 0

        # 检查是否能赢
        board[row][col] = -1
        if self.check_win(board, -1):
            board[row][col] = 0
            return 10000

        # 检查是否需要防守
        board[row][col] = 1
        if self.check_win(board, 1):
            board[row][col] = 0
            return 9000

        board[row][col] = 0

        # 评估进攻和防守潜力
        for dr, dc in self.directions:
            # 进攻评分
            attack_score = self.count_line(board, row, col, dr, dc, -1)
            score += attack_score * 10

            # 防守评分
            defense_score = self.count_line(board, row, col, dr, dc, 1)
            score += defense_score * 9

        # 中心位置加分
        center = self.board_size // 2
        distance_to_center = abs(row - center) + abs(col - center)
        score += (self.board_size - distance_to_center)

        # 根据难度调整评分权重
        if self.difficulty >= 4:
            # 高难度AI更重视连子数
            score *= 1.2
        elif self.difficulty >= 5:
            # 大师级AI更重视位置控制
            score *= 1.5

        return score

    def count_line(self, board, row, col, dr, dc, player):
        """计算某个方向上的连子数"""
        count = 1
        open_ends = 0

        # 正向检查
        r, c = row + dr, col + dc
        while 0 <= r < self.board_size and 0 <= c < self.board_size:
            if board[r][c] == player:
                count += 1
            elif board[r][c] == 0:
                open_ends += 1
                break
            else:
                break
            r, c = r + dr, c + dc

        # 反向检查
        r, c = row - dr, col - dc
        while 0 <= r < self.board_size and 0 <= c < self.board_size:
            if board[r][c] == player:
                count += 1
            elif board[r][c] == 0:
                open_ends += 1
                break
            else:
                break
            r, c = r - dr, c - dc

        # 根据连子数和开放端数返回分数
        if count >= 5:
            return 1000
        elif count == 4:
            if open_ends == 2:
                return 500
            elif open_ends == 1:
                return 100
        elif count == 3:
            if open_ends == 2:
                return 50
            elif open_ends == 1:
                return 10
        elif count == 2:
            if open_ends == 2:
                return 5
            elif open_ends == 1:
                return 2

        return 1

    def check_win(self, board, player):
        """检查是否获胜"""
        # 检查横向
        for i in range(self.board_size):
            for j in range(self.board_size - 4):
                if all(board[i][j + k] == player for k in range(5)):
                    return True

        # 检查纵向
        for i in range(self.board_size - 4):
            for j in range(self.board_size):
                if all(board[i + k][j] == player for k in range(5)):
                    return True

        # 检查主对角线
        for i in range(self.board_size - 4):
            for j in range(self.board_size - 4):
                if all(board[i + k][j + k] == player for k in range(5)):
                    return True

        # 检查副对角线
        for i in range(self.board_size - 4):
            for j in range(4, self.board_size):
                if all(board[i + k][j - k] == player for k in range(5)):
                    return True

        return False
